#include "safetensors_loader.h"
#include "macro.h"
#include "status.h"

#include <fstream>
#include <filesystem>
#include <unistd.h>
#include <sys/mman.h>
#include <fcntl.h> 

namespace iangels {

Ret SafeTensorsLoader::load(const std::vector<std::string>& filenames) {
    for (const auto& filename : filenames) {
        SafeTensor tensor;
        tensor.filename_ = filename;

        auto file_bytes = std::filesystem::file_size(filename);
        IANGELS_CHECK(file_bytes > sizeof(u_int64_t), FILE_SIZE_INVALID,
            "File size is invalid: %s", filename.c_str());
        tensor.file_bytes_ = file_bytes;

        std::ifstream file(filename, std::ifstream::binary);
        IANGELS_CHECK(file.is_open(), FILE_OPEN_FAILED,
            "Failed to open file: %s", filename.c_str());

        Ret ret = load_header_size(file, tensor);
        if (ret != SUCCESS) {
            LOG_ERROR("Failed to load header size: %s", filename.c_str());
            file.close();
            return ret;
        }

        ret = load_header(file, tensor);
        if (ret != SUCCESS) {
            LOG_ERROR("Failed to load header: %s", filename.c_str());
            file.close();
            return ret;
        }

        ret = mmap_weights(filename, tensor);
        if (ret != SUCCESS) {
            LOG_ERROR("Failed to mmap weights: %s", filename.c_str());
            file.close();
            return ret;
        }

        file.close();
        tensors_.emplace_back(std::move(tensor));
    }
    return SUCCESS;
}

Ret SafeTensorsLoader::load_header_size(std::ifstream& file, SafeTensor& tensor)
{
    file.read(reinterpret_cast<char*>(&(tensor.header_bytes_)), sizeof(uint64_t));
    auto bytes_read = file.gcount();
    IANGELS_CHECK(bytes_read == sizeof(uint64_t), FILE_READ_FAILED);

    // safetensors N 使用 little-endian
    // 检查系统是否使用 little-endian
    // TODO

    LOG_DEBUG("Get Header size: %ld\n", tensor.header_bytes_);
    return SUCCESS;
}

Ret SafeTensorsLoader::load_header(std::ifstream& file, SafeTensor& tensor)
{
    IANGELS_CHECK(tensor.file_bytes_ >= (sizeof(u_int64_t) + tensor.header_bytes_), FILE_SIZE_INVALID,
            "File size is invalid");

    if (tensor.header_bytes_ > MAX_HEADER_SIZE) {
        return FILE_SIZE_INVALID;
    }
    std::vector<char> buffer(tensor.header_bytes_);
    file.read(buffer.data(), tensor.header_bytes_);

    uint64_t bytes_read = file.gcount();
    IANGELS_CHECK(bytes_read == tensor.header_bytes_, FILE_READ_FAILED,
        "Failed to read header, expect: %ld, actual: %ld", tensor.header_bytes_, bytes_read);
    std::string header_str(buffer.begin(), buffer.end());
    tensor.header_ = Config::parse(header_str);
    
    return SUCCESS;
}

Ret SafeTensorsLoader::mmap_weights(const std::string& filename, SafeTensorsLoader::SafeTensor& tensor)
{
    int fd = open(filename.c_str(), O_RDWR);
    IANGELS_CHECK(fd > 0, FILE_OPEN_FAILED, "Failed to open file: %s", filename.c_str());

    char* file_addr = (char*)mmap(NULL, tensor.file_bytes_, PROT_READ, MAP_SHARED, fd, 0);
    IANGELS_CHECK(file_addr != MAP_FAILED, FILE_MMAP_FAILED, "Failed to mmap file: %s", filename.c_str());
    
    // close fd after mmap success
    close(fd);
    fd = -1;
    tensor.weights_addr_ = file_addr + sizeof(uint64_t) + tensor.header_bytes_;
    return SUCCESS;
}

Ret SafeTensorsLoader::unmap_weights()
{
    for (auto& tensor : tensors_) {
        munmap(tensor.weights_addr_, tensor.file_bytes_);
    }
    return SUCCESS;
}

} // namespace iangels